from core.bytecodes import LOAD_CONST, ADD, PRINT, STOP, OpMap, CALL, SET_GLOBAL, LOAD_GLOBAL, disassemble_codes, \
    RETURN, SET_LOCAL, LOAD_LOCAL
from core.compiler import Compiler
from core.objects import ExecutorObject, Value, CallFrame, ObjectType, Function, IS_NUM, AS_NUM

STACK_MAX_SIZE = 1024


class Constants:
    pass


class VM:
    # 常量池
    # 字节码
    def __init__(self):
        self.constants: list[Value] = []
        self.globals: dict[Value, Value] = {}
        self.stack: list[Value] = []
        self.stack_top = 0
        self.stack_frames: list[CallFrame] = [CallFrame() for _ in range(STACK_MAX_SIZE)]
        self.stack_max_size = STACK_MAX_SIZE
        self.current_frame_index = 0
        self.locals = []

    def current_frame(self) -> CallFrame:
        return self.stack_frames[self.current_frame_index]

    def create_frame(self, fn: Function) -> CallFrame:
        self.current_frame_index += 1
        if self.current_frame_index >= STACK_MAX_SIZE:
            raise Exception('Stack overflow')
        self.stack_frames[self.current_frame_index].init(fn, self.stack_top)
        return self.stack_frames[self.current_frame_index]

    def free_frame(self):
        self.stack_frames[self.current_frame_index].free()
        self.current_frame_index -= 1

    def load_const(self, index: int) -> Value:
        return self.constants[index]

    def push(self, value: Value):
        if self.stack_top >= STACK_MAX_SIZE:
            raise Exception('Stack overflow')
        self.stack.append(value)
        self.stack_top += 1

    def pop(self) -> Value:
        if self.stack_top <= 0:
            raise Exception('Stack underflow')
        self.stack_top -= 1
        return self.stack.pop()

    def run(self):
        while True:
            instruction = self.read_byte()
            if instruction == STOP:
                break
            elif instruction == LOAD_CONST:
                index = self.read_byte()
                val = self.load_const(index)
                self.push(val)
            elif instruction == ADD:
                val1 = self.pop()
                val2 = self.pop()
                if IS_NUM(val1) and IS_NUM(val2):
                    res = AS_NUM(val1.val + val2.val)
                else:
                    raise Exception("Not a number to Add")
                self.push(res)
            elif instruction == PRINT:
                val = self.pop()
                print(val.val)
            elif instruction == CALL:
                func = self.pop()

                if func.obj_type != ObjectType.PROCESS:
                    raise Exception("Not a function")

                # 创建函数调用栈帧
                call_frame = self.create_frame(func)
               # 从最后一个参数
                for i in range(func.arity - 1, -1, -1):
                    call_frame.add_local(func.args[i], self.pop())
            elif instruction == RETURN:
                return_val = self.pop()
                self.free_frame()
                self.push(return_val)

            elif instruction == SET_GLOBAL:
                val = self.pop()
                var_name = self.pop()
                self.add_global(var_name, val)
            elif instruction == LOAD_GLOBAL:
                var_name = self.pop()
                self.push(self.get_global(var_name))
            elif instruction == SET_LOCAL:
                val = self.pop()
                var_name = self.pop()
                self.add_local(var_name, val)
            elif instruction == LOAD_LOCAL:
                var_name = self.pop()
                self.push(self.get_local(var_name))
            else:
                raise Exception(f"Unknown instruction {OpMap[instruction]}")

    def read_byte(self) -> int:
        return self.current_frame().read_byte()

    def interpret(self, exec_obj: ExecutorObject):
        main_func = Function(None, [], exec_obj.codes)
        self.current_frame().init(main_func, 0)
        self.constants = exec_obj.constants
        self.globals = exec_obj.globals
        disassemble_codes(self.current_frame().get_codes())

        for k, v in exec_obj.globals.items():
            if v.obj_type == ObjectType.PROCESS:
                print(f"func {k.val}")
                disassemble_codes(v.body)
        print("start scheme vm...")
        self.run()

    def start(self, codes):
        compiler = Compiler()
        exec_obj = compiler.compile(codes)

        self.interpret(exec_obj)

    def add_global(self, var_name, val):
        if var_name in self.globals:
            raise Exception("Redefine global variable")
        else:
            self.globals[var_name] = val

    def get_global(self, var_name: Value):
        if var_name in self.globals:
            return self.globals[var_name]
        else:
            raise Exception("Undefined global variable")

    def add_local(self, var_name, val):
        self.current_frame().add_local(var_name, val)

    def get_local(self, var_name):
        for i in range(self.current_frame_index, -1, -1):
            if self.stack_frames[i].locals.get(var_name):
                return self.stack_frames[i].locals[var_name]


if __name__ == '__main__':
    codes = """
        (define y 2)
        (define (func1 x)
            (+ x y)
        )
        (print (func1 3))
    """
    vm = VM()
    vm.start(codes)
